import cv2
import numpy as np
import matplotlib.pyplot as plt


def gray_gamma(image_path, save_plot_path, gamma=1.0, save_path=None, show_result=True):
    # 读取图像
    img = cv2.imread(image_path, 0)
    if img is None:
        raise ValueError("无法读取图像，请检查路径")

    # 归一化到0-1范围
    normalized_img = img / 255.0
    # 应用伽马校正
    corrected_img = np.power(normalized_img, gamma)
    # 转换回0-255范围
    corrected_img = np.clip(corrected_img * 255.0, 0, 255).astype(np.uint8)

    # 结果展示与保存
    if show_result:
        plt.rcParams['font.sans-serif'] = ['Microsoft YaHei'] 
        plt.figure(figsize=(12, 10))
        plt.subplot(221), plt.imshow(img, cmap='gray'), plt.title('原始图像')
        plt.subplot(222), plt.imshow(corrected_img, cmap='gray'), plt.title(f'伽马校正后图像 (gamma={gamma})')
        # 绘制原始图像的直方图
        plt.subplot(223)
        plt.hist(img.ravel(), 256, [0, 256], color='b', alpha=0.7)
        plt.title('原始图像直方图')
        plt.xlabel('像素值')
        plt.ylabel('频率')
        plt.xlim([0, 256])
        
        # 绘制校正后图像的直方图
        plt.subplot(224)
        plt.hist(corrected_img.ravel(), 256, [0, 256], color='r', alpha=0.7)
        plt.title('伽马校正后直方图')
        plt.xlabel('像素值')
        plt.ylabel('频率')
        plt.xlim([0, 256])
        
        plt.tight_layout()  # 自动调整子图参数，使之填充整个图像区域
        plt.savefig(save_plot_path)
        plt.show()
    
    if save_path:
        cv2.imwrite(save_path, corrected_img)
    
    return corrected_img


if __name__ == "__main__":
    # 输入图像路径（请替换为你的图像路径）
    image_path = "./gray_gamma/original_3.jpg"
    # 输出路径（可选）
    save_path = "./gray_gamma/corrected_image_3.png"
    compare_path = "./gray_gamma/compare_image_3.png"

    # 执行灰度矫正
    corrected_img = gray_gamma(image_path, compare_path, gamma=2.2, save_path=save_path)